import asyncio
import json
import contextlib
import logging
from typing import Optional

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, HTTPException
from sqlmodel import Session
from datetime import datetime, timezone

from ..config import get_settings
from ..db import engine
from ..models import Server, SessionLog
from ..services.crypto import decrypt_text
from ..services.ssh import SSHShellSession
from ..services.settings import read_settings


router = APIRouter()
logger = logging.getLogger(__name__)


@router.websocket("/ws/ssh/{server_id}")
async def websocket_ssh(websocket: WebSocket, server_id: int, cols: Optional[int] = 120, rows: Optional[int] = 32):
    await websocket.accept()
    try:
        client_info = getattr(websocket, "client", None)
        logger.info(
            "[WS] accepted server_id=%s from=%s",
            server_id,
            f"{getattr(client_info, 'host', 'unknown')}:{getattr(client_info, 'port', 'unknown')}" if client_info else "unknown",
        )
    except Exception:  # noqa: BLE001
        logger.debug("[WS] accepted (client info unavailable)")
    # 提前告知前端通路已建立，可进行初始交互（将缓冲到 SSH 连接完成）
    with contextlib.suppress(Exception):
        await websocket.send_text(json.dumps({"type": "accepted"}))

    # Load server info
    with Session(engine) as session:
        server = session.get(Server, server_id)
        if server is None:
            logger.error("[WS] server not found: id=%s", server_id)
            await websocket.send_text(json.dumps({"type": "error", "message": "Server not found"}))
            await websocket.close(code=4004)
            return

    settings = get_settings()
    runtime = read_settings()
    password = decrypt_text(server.password_enc or "")

    # Establish SSH session
    ssh = SSHShellSession(
        host=server.host,
        port=server.port,
        username=server.username,
        password=password or None,
        width=cols or 120,
        height=rows or 32,
        timeout=int(runtime.get("ssh_timeout_seconds", settings.ssh_timeout_seconds)),
    )

    try:
        logger.info(
            "[WS] SSH connect start host=%s port=%s username=%s timeout=%s cols=%s rows=%s",
            server.host,
            server.port,
            server.username,
            settings.ssh_timeout_seconds,
            cols,
            rows,
        )
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(None, ssh.connect)
        logger.info("[WS] SSH connect success host=%s port=%s username=%s", server.host, server.port, server.username)
    except Exception as exc:  # noqa: BLE001
        logger.exception("[WS] SSH connect failed host=%s port=%s username=%s", server.host, server.port, server.username)
        await websocket.send_text(json.dumps({"type": "error", "message": f"SSH connect failed: {exc}"}))
        await websocket.close(code=4005)
        return

    send_task = None
    log_buffer: list[str] = []
    bytes_count = 0
    connected_to_ssh = True
    buffered_inputs: list[str] = []
    pending_resize: Optional[tuple[int, int]] = None
    log_row_id: Optional[int] = None
    flushed_chars = 0

    async def pump_ssh_to_ws():
        nonlocal bytes_count
        nonlocal flushed_chars
        try:
            async for chunk in ssh.iterate_stdout():
                try:
                    text = chunk.decode("utf-8", errors="ignore")
                    log_buffer.append(text)
                    bytes_count += len(chunk)
                    if bytes_count % (64 * 1024) < 4096:  # 降低日志噪音：每累计约 64KB 打一条
                        logger.debug("[WS] pump stdout total_bytes=%s", bytes_count)
                    await websocket.send_text(json.dumps({"type": "data", "data": text}))
                    # 追写 bytes_count 与部分 content，便于“登录后立即可见内容片段与大小”
                    flush_threshold = int(runtime.get("log_flush_threshold_bytes", 16 * 1024))
                    if log_row_id is not None and bytes_count % flush_threshold < 4096:
                        with contextlib.suppress(Exception):
                            with Session(engine) as session:
                                row = session.get(SessionLog, log_row_id)
                                if row is not None:
                                    row.bytes_count = bytes_count
                                    # 将未刷新的部分追加写入（基于字符数）
                                    current_full = "".join(log_buffer)
                                    if flushed_chars < len(current_full):
                                        row.content = (row.content or "") + current_full[flushed_chars:]
                                        flushed_chars = len(current_full)
                                    session.add(row)
                                    session.commit()
                except Exception:  # noqa: BLE001
                    logger.exception("[WS] pump_ssh_to_ws error while sending to client")
                    break
        finally:
            # 不在这里主动关闭，由外层 finally 统一处理
            pass

    # 先创建会话日志占位（便于前端列表立即可见）
    try:
        with Session(engine) as session:
            row = SessionLog(
                server_id=server.id,
                started_at=datetime.now(timezone.utc),
                ended_at=None,
                bytes_count=0,
                content="",
            )
            session.add(row)
            session.commit()
            session.refresh(row)
            log_row_id = row.id
            logger.info("[WS] session log placeholder created id=%s server_id=%s", log_row_id, server.id)
    except Exception:
        logger.exception("[WS] failed to create session log placeholder server_id=%s", server.id)

    # 告知前端 SSH 就绪
    try:
        await websocket.send_text(json.dumps({"type": "ready"}))
        # 触发远端 shell 打印提示符
        try:
            await ssh.write_stdin(b"\n")
        except Exception:
            pass
        # 向前端输出一行欢迎信息，便于确认通路
        with contextlib.suppress(Exception):
            await websocket.send_text(json.dumps({"type": "data", "data": "\r\n[WebSSH 已连接，等待远端输出...]\r\n"}))
        logger.info("[WS] ready sent; start receiving messages server_id=%s", server_id)
    except Exception:
        logger.exception("[WS] failed to send ready; closing ssh")
        with contextlib.suppress(Exception):
            await ssh.close()
        return

    send_task = asyncio.create_task(pump_ssh_to_ws())

    try:
        while True:
            message = await websocket.receive_text()
            try:
                payload = json.loads(message)
            except Exception:
                # Treat as raw keystrokes
                if connected_to_ssh:
                    await ssh.write_stdin(message.encode("utf-8"))
                else:
                    buffered_inputs.append(message)
                continue

            msg_type = payload.get("type")
            if msg_type == "input":
                data = payload.get("data", "")
                logger.debug("[WS] recv input len=%s", len(data) if isinstance(data, str) else 0)
                if connected_to_ssh:
                    await ssh.write_stdin(data.encode("utf-8"))
                else:
                    buffered_inputs.append(str(data))
            elif msg_type == "resize":
                cols = int(payload.get("cols", 120))
                rows = int(payload.get("rows", 32))
                logger.info("[WS] recv resize cols=%s rows=%s", cols, rows)
                if connected_to_ssh:
                    await ssh.resize_pty(width=cols, height=rows)
                else:
                    pending_resize = (cols, rows)
            elif msg_type == "ping":
                logger.debug("[WS] recv ping; replying pong")
                await websocket.send_text(json.dumps({"type": "pong"}))
            else:
                # Ignore unknown message types
                logger.debug("[WS] recv unknown type: %s", msg_type)
                pass
    except WebSocketDisconnect:
        logger.info("[WS] client disconnected server_id=%s", server_id)
    except Exception as exc:  # noqa: BLE001
        logger.exception("[WS] error in message loop server_id=%s: %s", server_id, exc)
        try:
            await websocket.send_text(json.dumps({"type": "error", "message": str(exc)}))
        except Exception:
            pass
    finally:
        if send_task:
            send_task.cancel()
            with contextlib.suppress(Exception):
                await send_task
        await ssh.close()
        logger.info(
            "[WS] closing session server_id=%s bytes_count=%s log_chunks=%s",
            server_id,
            bytes_count,
            len(log_buffer),
        )
        # Persist session log（更新或补写）
        try:
            content = "".join(log_buffer)
            with Session(engine) as session:
                if log_row_id is not None:
                    row = session.get(SessionLog, log_row_id)
                    if row is not None:
                        row.ended_at = datetime.now(timezone.utc)
                        row.bytes_count = bytes_count
                        # 以完整内容覆盖，确保与增量写入一致
                        row.content = content
                        session.add(row)
                        session.commit()
                        logger.info("[WS] session log updated id=%s size=%s", log_row_id, len(content))
                    else:
                        # 回退创建
                        backup = SessionLog(
                            server_id=server.id,
                            started_at=datetime.now(timezone.utc),
                            ended_at=datetime.now(timezone.utc),
                            bytes_count=bytes_count,
                            content=content,
                        )
                        session.add(backup)
                        session.commit()
                        logger.info("[WS] session log inserted (fallback) server_id=%s size=%s", server_id, len(content))
                else:
                    # 未能创建占位时，直接插入
                    row = SessionLog(
                        server_id=server.id,
                        started_at=datetime.now(timezone.utc),
                        ended_at=datetime.now(timezone.utc),
                        bytes_count=bytes_count,
                        content=content,
                    )
                    session.add(row)
                    session.commit()
                    logger.info("[WS] session log inserted server_id=%s size=%s", server_id, len(content))
        except Exception:
            logger.exception("[WS] failed to persist session log server_id=%s", server_id)

